import sys, os, string, datetime, re
from Lib import Util
from ftplib import FTP

# Download all the RouteViews archives into the correct directories for
# a given span of time from start->end.
# Note: We don't take the minute in the argument, just the hour.
# I figure that getting a little extra data is ok, and in most cases
# the usage would be to add 00 as the minutes anyway.
#---------------------------------------#

# Given an datetime and an existing ftp connection
# and download all the update archives in the current remote directory 
# for that hour to the current local directory.
# Note: It's up to something else to make sure the current remote and local
# directories are where they're supposed to be.
def get_hour(ftp, dt):

        # Get the file listing only if we haven't previously retrieved it
        if not file_lists.has_key(dt):
            file_lists[dt] = ftp.nlst()

        # Parse out the files we want for this
        files = [f for f in file_lists[dt] if 'updates.%04d%02d%02d.%02d' % 
        (dt.year, dt.month, dt.day, dt.hour) in f]

        # Download all the files
        for f in files:

            # If the file exists and the sizes match, don't
            # bother downloading it again.
            # Otherwise, delete the file if necessary and download
            skip = 0
            if os.path.exists(f):
                
                # Files are the same size, skip
                if ftp.size(f) == os.path.getsize(f):
                    skip = 1
                    print '%s exists, skipping.' % f

                # Files are not the same, delete the local copy 
                # and reset the skip flag
                else:
                    os.unlink(f)
                    skip = 0
            
            if f[-1] == '~':
                skip = 1
                
            # We want the file
            if not skip:
                print 'Retrieving %s' % f
                try:
                    # Retrieve and write to the current directory
                    ftp.retrbinary('RETR %s' % f, open(f, 'w', 0664).write)
                except:
                    print 'Error fetching file %s' % f
                    try:
                        # Abort raises an exception if it's not currently
                        # transferring data
                        ftp.abort()
                        ftp.quit()
                    except:
                        ftp.quit()


# Establish an ftp connection and do the directory management
# as we loop through time and try to get all the archives
def get_archives(monitor, start_dt, end_dt):

    # RIPE archive, use HTTP
    if 'rrc' in monitor:
        get_html_archives(monitor, start_dt, end_dt)
        
    # RV archive, use FTP
    else:
        get_ftp_archives(monitor, start_dt, end_dt)
        
        
        
def get_html_archives(monitor, start_dt, end_dt):
    current_dt = start_dt
    current_month = 0
    while(current_dt <= end_dt):
        current_month = current_dt.month
        # Regexp for matching update archive links in the html page
        re_line = re.compile('(updates\.%04d%02d\d\d.\d\d\d\d\.gz)' % (current_dt.year, current_dt.month))

        # Local directory structure
        local_dir = '%s/%04d.%02d' % \
                (monitor, int(current_dt.year), int(current_dt.month))

        # If we don't have the directory locally, create it.
        if not os.path.exists('%s/%s' % (Util.archive_dir, local_dir)):
            os.makedirs('%s/%s' % (Util.archive_dir, local_dir), 0775)
        if not os.path.exists('%s/%s/mrt' % (Util.archive_dir, local_dir)):
            os.makedirs('%s/%s/mrt' % (Util.archive_dir, local_dir), 0775)
        if not os.path.exists('%s/%s/plain' % (Util.archive_dir, local_dir)):
            os.makedirs('%s/%s/plain' % (Util.archive_dir, local_dir), 0775)

        # Change the local dir to where we want the files to go
        os.chdir('%s/%s/mrt' % (Util.archive_dir, local_dir))

        # No file list exists
        if not os.path.isfile('index.html'):
            os.system('%s %s/%s/%04d.%02d' % (wget, http_host, monitor, current_dt.year, current_dt.month))
        
        # We have data for the entire month
        while(current_dt.month == current_month and current_dt <= end_dt):
            # We should have an index file if any archives exist for this date
            if os.path.isfile('index.html'):
                fh = open('index.html', 'r')
                line = string.rstrip(fh.readline())
                fh_list = open('file_list', 'w', 0664)  # Processed links for wget
                while(line):
                    match = re_line.search(line)
                    if match:   # Matches our archive file format, add to the list
                        if 'updates.%04d%02d%02d.%02d' % (current_dt.year, current_dt.month, current_dt.day, current_dt.hour) in line:
                            fh_list.write('%s/%s/%04d.%02d/%s\n' % (http_host, monitor, current_dt.year, current_dt.month, match.group(1)))
                    line = fh.readline()
                fh_list.close()
                fh.close()
       
            # If there are archives to download, have wget grab them
            if os.path.isfile('file_list'):
                os.system('%s -i %s' % (wget, 'file_list'))

            # Go to the next hour in the month
            current_dt = current_dt + datetime.timedelta(hours=1)
    



def get_ftp_archives(monitor, start_dt, end_dt):

    # Login to the host
    ftp = FTP(ftp_host)
    ftp.login()
   
    # Clear our lists because this is a different monitor
    file_lists.clear() 
   
    # Loop through time incrementing by hour
    current_dt = start_dt
    while(current_dt <= end_dt):
        # Local directory structure
        local_dir = '%s/%04d.%02d' %\
            (monitor, int(current_dt.year), int(current_dt.month))
       
        # route-views2 and oix-route-views has a special directory structure
        # Everything else uses the intermediate bgpdata directory
        if monitor == 'route-views2':
            remote_dir = '/bgpdata/%04d.%02d' % (int(current_dt.year), int(current_dt.month))
        else:
            raise #this needs to be fixed, the directory structure isn't like this anymore
            #remote_dir = '%s/bgpdata/%04d.%02d' % \
            #(monitor, int(current_dt.year), int(current_dt.month))

        # If we don't have the directory locally, create it.
        if not os.path.exists('%s/%s' % (Util.archive_dir, local_dir)):
            os.makedirs('%s/%s' % (Util.archive_dir, local_dir), 0775)
        if not os.path.exists('%s/%s/mrt' % (Util.archive_dir, local_dir)):
            os.makedirs('%s/%s/mrt' % (Util.archive_dir, local_dir), 0775)
        if not os.path.exists('%s/%s/plain' % (Util.archive_dir, local_dir)):
            os.makedirs('%s/%s/plain' % (Util.archive_dir, local_dir), 0775)

        # Change the local dir to where we want the files to go
        os.chdir('%s/%s/mrt' % (Util.archive_dir, local_dir))

        # Change the remote dir and get the archives
        temp_remote = '%s/UPDATES' % remote_dir

        try:
            ftp.cwd(temp_remote)
        except:
            # The directory we want doesn't exist for that monitor.
            print "Could not change to %s" % temp_remote

        # Get all the archives for the hour
        get_hour(ftp, current_dt)

        # Update the time + 1 hour
        current_dt = current_dt + datetime.timedelta(hours=1)
    
    # We got all the archives for this monitor
    ftp.quit()


#----Main----#

ftp_host = 'archive.routeviews.org'
http_host = 'http://data.ris.ripe.net'
wget = '/usr/bin/wget -N' # Check timestamp, quiet
assert os.path.exists(Util.archive_dir)

file_lists = {}

if len(sys.argv) < 3:
    print "Usage: fetch_archives.py YYYY-MM-DD YYYY-MM-DD [monitor]"
    sys.exit(-1)

# So we can set 775/664 permissions on things
os.umask(002)

# Turn the start and end arguments into datetimes
start = string.split(sys.argv[1], '-')
end = string.split(sys.argv[2], '-')
start_dt = datetime.datetime(int(start[0]), int(start[1]), 
        int(start[2]), 0, 0, 0)
end_dt = datetime.datetime(int(end[0]), int(end[1]), 
        int(end[2]), 23, 59, 59)
monitors = sys.argv[3:]
if not monitors:
    monitors = Util.monitors
if monitors[0] == 'rv':
    monitors = Util.rv_monitors
elif monitors[0] == 'ripe':
    monitors = Util.ripe_monitors

print monitors
# For each RouteViews monitors download the archives we want
for monitor in monitors:
    print 'Downloading updates from monitor %s' % monitor
    get_archives(monitor, start_dt, end_dt)